package cc.shacocloud.kotlin.tools.http

import okhttp3.OkHttpClient
import java.io.IOException
import java.io.InputStream
import java.security.GeneralSecurityException
import java.security.KeyStore
import java.security.cert.CertificateFactory
import javax.net.ssl.*

/**
 * 自定义TLS证书握手工具
 *
 * @author 思追(shaco)
 */
object CustomizeTLSCertificateHandshakeUtil {

    /**
     * 构建一个新的自签TLS握手的 OkHttpClient
     *
     * @param client 旧的 OkHttpClient
     */
    @Throws(RuntimeException::class)
    fun builder(
        client: OkHttpClient,
        certificate: InputStream
    ): OkHttpClient {
        val trustManager: X509TrustManager
        val socketFactory: SSLSocketFactory
        try {
            trustManager = trustManagerForCertificates(certificate)
            val sslContext = SSLContext.getInstance("TLS")
            sslContext.init(null, arrayOf<TrustManager>(trustManager), null)
            socketFactory = sslContext.socketFactory
        } catch (e: Exception) {
            throw RuntimeException(e)
        }

        // 构建 OkHttpClient
        return client.newBuilder()
            .sslSocketFactory(socketFactory, trustManager)
            .build()
    }

    /**
     * 以流的方式添加信任证书
     */
    @Throws(GeneralSecurityException::class, IOException::class)
    private fun trustManagerForCertificates(certificateInputStream: InputStream): X509TrustManager {
        val certificateFactory = CertificateFactory.getInstance("X.509")
        val certificates = certificateFactory.generateCertificates(certificateInputStream)
        require(!certificates.isEmpty()) { "受信任证书集为空！" }

        // 将证书作为密钥存储。
        // 任何密码都可以使用。
        val password = "password".toCharArray()
        val keyStore = newEmptyKeyStore(password)
        var index = 0
        for (certificate in certificates) {
            val certificateAlias = index++.toString()
            keyStore.setCertificateEntry(certificateAlias, certificate)
        }

        // 使用它来构建 X509 信任管理器。
        val keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm())
        keyManagerFactory.init(keyStore, password)

        val trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
        trustManagerFactory.init(keyStore)

        val trustManagers = trustManagerFactory.trustManagers
        check(!(trustManagers.size != 1 || trustManagers[0] !is X509TrustManager)) { "意外的默认信任管理器：" + trustManagers.contentToString() }
        return trustManagers[0] as X509TrustManager
    }


    /**
     * 添加 password
     */
    @Throws(GeneralSecurityException::class, IOException::class)
    private fun newEmptyKeyStore(password: CharArray): KeyStore {
        val keyStore = KeyStore.getInstance(KeyStore.getDefaultType())
        // 这里添加自定义的密码，默认
        // 按照惯例，“null” 会创建一个空密钥存储。
        keyStore.load(null, password)
        return keyStore
    }
}
